import dotenv
import weaviate
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.language_models import BaseLanguageModel
from langchain_community.embeddings.baidu_qianfan_endpoint import QianfanEmbeddingsEndpoint
from langchain_core.retrievers import BaseRetriever
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langchain_weaviate import WeaviateVectorStore
from langchain_core.runnables import RunnablePassthrough
from langchain_core.output_parsers import StrOutputParser
from weaviate.auth import AuthApiKey

dotenv.load_dotenv()

class HyDERectriever(BaseRetriever):
    """HyDE混合策略检索器"""
    retriever: BaseRetriever
    llm: BaseLanguageModel

    def _get_relevant_documents(
        self, query: str, *, run_manager: CallbackManagerForRetrieverRun
    ) -> list[Document]:
        """传递检索query实现HyDE混合策略检索"""
        # 1. 构建生成假设性文档的prompt
        prompt = ChatPromptTemplate.from_template(
            "请写一篇短的科学论文来发这个问题。\n"
            "问题：{question}\n"
            "文章："
        )

        # 2. 构建HyDE混合策略检索链
        chain = (
            {"question": RunnablePassthrough()}
            | prompt
            | self.llm
            | StrOutputParser()
            | self.retriever
        )

        return chain.invoke(query)

# 3. 构建向量数据库与检索器
# 创建客户端连接（使用新的connect_to_weaviate_cloud方法）
client = weaviate.connect_to_weaviate_cloud(
    cluster_url="https://zabwh0mbt4errmvpknamq.c0.asia-southeast1.gcp.weaviate.cloud",
    auth_credentials=AuthApiKey("b2o4OGQxcmptMTZEWmJ5VV9udE5xSXBzQW04dUlDZ0JSS0d1ay9FQlhXdEtyMDR4OUFVNzc0eG9mU3dnPV92MjAw")
)

db = WeaviateVectorStore(
    client=client,
    index_name="myleane",
    text_key="text",
    embedding=QianfanEmbeddingsEndpoint(),
)
retriever = db.as_retriever(search_type="mmr")

# 3. 构建HyDE检索器
hyde_retriever = HyDERectriever(
    retriever = retriever,
    llm = ChatOpenAI(model_name="kimi-k2-0711-preview", temperature=0)
)


# 4. 执行检索
docs = hyde_retriever.invoke("关于LLMOps应用配置的文档有哪些")
client.close()
print(docs)
print(len(docs))
